import os
import re
import email
import csv
import argparse
from tqdm import tqdm
import chardet
from bs4 import BeautifulSoup
import traceback
import sys

class EmailProcessor:
    def __init__(self, input_dir, csv_path="emails.csv", batch_size=100, min_valid_chars=100):
        self.input_dir = input_dir
        self.csv_path = csv_path
        self.batch_size = batch_size
        self.min_valid_chars = min_valid_chars  # 最小有效字符数
        self.emails_batch = []
        self._init_csv()

    def _init_csv(self):
        with open(self.csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile, escapechar='\\', quoting=csv.QUOTE_ALL)
            writer.writerow([
                'message_id', 'file_path', 'from_email', 'subject', 
                'date', 'body', 'labels', 'spam_label', 'status'
            ])

    def _force_decode(self, payload):
        if isinstance(payload, email.message.Message):
            charset = payload.get_content_charset()
            if charset:
                try:
                    decoded = payload.get_payload(decode=True).decode(charset, errors='replace')
                    return decoded, charset, True  
                except:
                    pass

 
        sample_size = min(len(payload), 10000) 
        result = chardet.detect(payload[:sample_size])
        encoding = result['encoding']
        confidence = result['confidence']

        if confidence > 0.7 and encoding in ['utf-8', 'gb2312', 'gbk', 'big5', 'iso-8859-1', 'latin-1']:
            try:
                decoded = payload.decode(encoding, errors='replace')
                return decoded, encoding, confidence > 0.9 
            except:
                pass

       
        common_encodings = [
            'utf-8', 'gb18030', 'gbk', 'big5',  # 中文字符集
            'iso-8859-1', 'latin-1', 'ascii',    # 西文字符集
            'utf-16', 'utf-32',                  # 其他Unicode变体
        ]

        for enc in common_encodings:
            try:
                decoded = payload.decode(enc, errors='replace')
                # 检查解码后是否包含足够的有效字符
                valid_chars = sum(1 for c in decoded if c.isprintable() or c.isspace())
                if valid_chars > self.min_valid_chars / 2:  # 至少有一半的最小有效字符
                    return decoded, enc, False  # 标记为不确定解码
            except:
                continue

        # 完全无法解码，返回None
        return None, None, False

    def _is_valid_body(self, body):
        if not body:
            return False

        # 过滤解码失败的占位符
        if body.startswith("[无法解码") or body.startswith("[Error"):
            return False

        # 过滤Unicode乱码（如FFFB系列）
        if re.search(r'\\uFFFB|�{3,}', body):  # �是Unicode替换字符，FFFB是常见乱码
            return False

        # 过滤特殊符号密集内容（非文字字符占比过高）
        text_chars = sum(1 for c in body if c.isalnum() or c in '，。！？：；“”‘’、')
        total_chars = len(body)
        if total_chars == 0:
            return False
        if text_chars / total_chars < 0.6:  # 文字字符占比需>60%
            return False

        # 过滤有效字符不足的内容
        valid_chars = sum(1 for c in body if c.isprintable() or c.isspace())
        if valid_chars < self.min_valid_chars:
            return False

        return True

    def _extract_body(self, msg):
        body_parts = []
        valid_parts = False

        try:
            if msg.is_multipart():
                for part in msg.walk():
                    content_type = part.get_content_type()
                    payload = part.get_payload(decode=True)

                    if not payload or not isinstance(payload, bytes):
                        continue

                    decoded, encoding, is_reliable = self._force_decode(payload)

                    if decoded:
                        # 处理HTML内容
                        if content_type == "text/html":
                            try:
                                from html import unescape
                                decoded = unescape(decoded)
                                soup = BeautifulSoup(decoded, "html.parser")
                                text = soup.get_text(separator='\n', strip=True)
                                body_parts.append(text)
                            except:
                                body_parts.append(decoded)  # 解析失败时保留原始文本
                        else:
                            body_parts.append(decoded)

                        # 标记有效部分
                        if is_reliable or self._is_valid_body(decoded):
                            valid_parts = True
            else:
                payload = msg.get_payload(decode=True)
                if payload and isinstance(payload, bytes):
                    decoded, encoding, is_reliable = self._force_decode(payload)
                    if decoded:
                        if msg.get_content_type() == "text/html":
                            try:
                                from html import unescape
                                decoded = unescape(decoded)
                                soup = BeautifulSoup(decoded, "html.parser")
                                body_parts.append(soup.get_text(separator='\n', strip=True))
                            except:
                                body_parts.append(decoded)
                        else:
                            body_parts.append(decoded)

                        if is_reliable or self._is_valid_body(decoded):
                            valid_parts = True

        except Exception as e:
            print(f"提取邮件正文时发生严重错误: {e}")
            traceback.print_exc()

        # 合并所有部分
        body = "\n\n".join(body_parts).strip()
        # 清理控制字符，但保留换行符和制表符
        body = re.sub(r'[\x00-\x08\x0B-\x0C\x0E-\x1F\x7F]', '', body)

        # 最终有效性检查
        return body if valid_parts and self._is_valid_body(body) else None

    def _process_email(self, file_path):
        """处理单个邮件，只返回有效内容"""
        try:
            with open(file_path, 'rb') as f:
                msg_data = f.read()

            msg = email.message_from_bytes(msg_data)

            from_email = msg.get('From', 'unknown')
            subject = msg.get('Subject', 'no_subject')
            date = msg.get('Date', 'unknown')

            body = self._extract_body(msg)

            # 如果无法提取有效正文，返回None
            if not body:
                return None

            labels = self._extract_labels(subject)
            labels_str = ",".join(labels) if labels else ""

            spam_label = "unknown"
            if "spam" in file_path.lower():
                spam_label = "spam"
            elif "ham" in file_path.lower():
                spam_label = "ham"

            return {
                'message_id': os.path.basename(file_path),
                'file_path': file_path,
                'from_email': from_email,
                'subject': subject,
                'date': date,
                'body': body,
                'labels': labels_str,
                'spam_label': spam_label,
                'status': "processed"
            }

        except Exception as e:
            print(f"处理文件 {file_path} 时出错: {e}")
            return None

    def _extract_labels(self, subject):
        """从主题提取标签"""
        if not subject:
            return []
        try:
            return [label.strip().lower() for label in re.findall(r'\[(.*?)\]', subject) if label.strip()]
        except:
            return []

    def process_emails(self):
        """处理所有邮件，只保存有效结果"""
        all_files = []
        for root, _, files in os.walk(self.input_dir):
            for file in files:
                all_files.append(os.path.join(root, file))

        print(f"找到 {len(all_files)} 个文件")
        processed_count = 0
        skipped_count = 0

        for file_path in tqdm(all_files, desc="处理邮件"):
            email_data = self._process_email(file_path)

            if email_data:
                processed_count += 1
                self.emails_batch.append(email_data)

                if len(self.emails_batch) >= self.batch_size:
                    self._save_batch()
            else:
                skipped_count += 1

        # 保存剩余批次
        if self.emails_batch:
            self._save_batch()

        print(f"处理完成:")
        print(f"  - 成功: {processed_count}")
        print(f"  - 跳过: {skipped_count}")
        print(f"结果已保存到 {self.csv_path}")

    def _save_batch(self):
        """保存批次数据到CSV"""
        if not self.emails_batch:
            return

        with open(self.csv_path, 'a', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile, escapechar='\\', quoting=csv.QUOTE_ALL)
            for entry in self.emails_batch:
                writer.writerow([
                    entry['message_id'],
                    entry['file_path'],
                    entry['from_email'],
                    entry['subject'],
                    entry['date'],
                    entry['body'],
                    entry['labels'],
                    entry['spam_label'],
                    entry['status']
                ])

        self.emails_batch = []

def main():
    parser = argparse.ArgumentParser(description="处理邮件并保存为CSV，过滤乱码和无效内容")
    parser.add_argument("--dataset", default="trec05p-1/data", help="数据集目录路径")
    parser.add_argument("--csv", default="clean_emails.csv", help="输出CSV文件")
    parser.add_argument("--batch", type=int, default=100, help="批次大小")
    parser.add_argument("--min-chars", type=int, default=100, help="最小有效字符数")

    args = parser.parse_args()

    if not os.path.exists(args.dataset):
        print(f"错误: 数据集目录 '{args.dataset}' 不存在")
        sys.exit(1)

    # 检查依赖
    required_packages = {
        'chardet': 'chardet',
        'bs4': 'beautifulsoup4',
        'tqdm': 'tqdm'
    }

    for imp, pkg in required_packages.items():
        try:
            __import__(imp)
        except ImportError:
            print(f"安装依赖: {pkg}")
            os.system(f"{sys.executable} -m pip install {pkg}")

    processor = EmailProcessor(
        input_dir=args.dataset,
        csv_path=args.csv,
        batch_size=args.batch,
        min_valid_chars=args.min_chars
    )

    processor.process_emails()

if __name__ == "__main__":
    main()
